#sort_utils.py
import parse_utils

#contains predicates mapped to sorts
class SortDB:
	data = {}
	def __init__(self, P):
		for (l,rdec) in P:
			try:
				predec = parse_utils.getPredicate(rdec)
			except Exception as e:
				raise Exception('Syntax error at line '+str(l))
			arity = len(predec[1])
			self.data[predec[0]+str(arity)] = predec[1]
	def getSort(self, pred, arity, pos):
		dec = self.data[pred+str(arity)]
		return dec[pos]

#sort decorations
#list of tuples (E,I,C) with sort constraint atoms
class SortDec:
	def __init__(self, rule, line, sortdb):
		self.hdec = []
		self.bdec = []
		(type, head, body) = rule
		self.sortdb = sortdb
		for constr in head:
			if constr[0] == 3:
				self.hdec.append(self.getAtomDec(constr,line))
			elif constr[0] == 2:
				self.hdec.append(self.getAggDec(constr,line))
			elif constr[0] == 1:
				self.hdec.append(self.getAtomDec(constr,line))
			else:
				self.hdec.append((None,None,None))
		for constr in body:
			if constr[0] == 3:
				self.bdec.append(self.getAtomDec(constr,line))
			elif constr[0] == 2:
				self.bdec.append(self.getAggDec(constr,line))
			elif constr[0] == 1:
				self.bdec.append(self.getAtomDec(constr,line))
			else:
				self.bdec.append((None,None,None))
	def getAtomDec(self, atom, line):
		predicate = atom[1]
		E = []
		C = []
		#for each term
		i = 0
		arity = len(atom[3])
		while i<arity:
			term = atom[3][i]
			try:
				sort = self.sortdb.getSort(predicate,arity,i)
			except IndexError:
				raise Exception('Predicate \''+predicate+'\' declared with arity less than '+str(i+1))
			except KeyError:
				raise Exception('Predicate \''+predicate+'\' with arity '+str(arity)+' undeclared')
			vars = parse_utils.getVars(term)
			atom = sort+'('+term+')'
			#add to decoration
			if not term in E:
				if len(vars)>0:
					E.append(atom)
				else:
					C.append(atom)
			i += 1
		return (E,None,C)
	def getAggDec(self, agg, line):
		return (None,None,None)
	def getQuantDec(self, agg, line):
		return (None,None,None)